"""
CSSIM智能体环境基础类 - 带详细注释的版本
作者：CSSIM开发团队
功能：提供CSSIM(Combat Simulation System)智能体环境的基础框架
包含智能体管理、环境重置、步骤执行、观测处理等核心功能
"""

# 导入标准库
import json, copy, re, os, inspect, os  # JSON处理、深拷贝、正则表达式、系统操作、反射
import time  # 时间相关操作
from time import sleep  # 延时函数

# 导入科学计算库
import numpy as np  # 数值计算
from Utils.tensor_ops import my_view, repeat_at  # 张量操作工具
from Common.base_env import RawObsArray  # 原始观测数组类

# 导入CSSIM相关模块
from Common.agent import Agent  # 智能体类
from Common.cssim_env_wrapper import CssimEnv, ScenarioConfig  # CSSIM环境包装器和场景配置
from .cython_env_func import tear_num_arr  # Cython优化的数字数组处理函数
from Common.action_set_encoder import digit2act_dictionary, AgentPropertyDefaults  # 动作编码和智能体属性默认值
from Common.action_set_encoder import decode_action_as_string, digitsToStrAction  # 动作解码函数
from config import GlobalConfig  # 全局配置
from Common.route_and_environment import load_ScenarioConfig  # 场景配置加载

# ============================================================================
# CSSIM通用功能类 - 继承自CssimEnv
# ============================================================================

class CssimCommonFn(CssimEnv):
    """
    CSSIM通用功能类
    这是CSSIM仿真环境的核心类，负责：
    1. 智能体生命周期管理（创建、重置、更新）
    2. 环境状态管理（episode、步骤、配置）
    3. 与UE（Unreal Engine）的通信
    4. 观测数据处理和格式化
    5. 奖励计算和任务结束判断
    """

    def __init__(self, rank) -> None:
        """
        初始化CSSIM环境

        Args:
            rank: 进程/线程的排名标识
        """
        super().__init__(rank)  # 调用父类CssimEnv的初始化

        # 基础状态变量
        self.train_mode = None          # 训练模式标志
        self.current_step = None        # 当前仿真步骤
        self.key_obj = None             # 关键游戏对象

        # 智能体相关变量
        self.num_agent_each_team = None  # 每个队伍的智能体数量
        self.n_team_agent = None         # 各队伍智能体数量字典
        self.episode = 0                 # 当前episode计数
        self.max_episode = None          # 最大episode数
        self.agents = []                 # 智能体对象列表

        # 目标位置
        self.target_position = None

        # 环境配置
        self.observation_space = self.make_obs(get_shape=True)  # 观测空间维度

        # 从全局配置中获取子任务配置
        self.SubTaskConfig = GlobalConfig.ScenarioConfig.SubTaskConfig

        # 奖励分享机制：环境为整个团队提供奖励，而不是单个智能体
        self.Enable_Reward_Sharing = False

        # 智能体设置数组
        self.AgentSettingArray = []  # 活跃智能体的设置（排除指挥官和中立方）
        self.AgentSettingAll = []    # 所有智能体的设置（包括指挥官和中立方）
        self.uid_to_agent_dict = {}  # UID到智能体的映射字典（缓存机制）

        # 初始化智能体列表
        self.SubTaskConfig.agent_list = self.AgentSettingArray.copy()

        # 检查是否启用奖励分享机制
        if hasattr(GlobalConfig.ScenarioConfig, 'Enable_Reward_Sharing'):
            self.Enable_Reward_Sharing = GlobalConfig.ScenarioConfig.Enable_Reward_Sharing

    # ========================================================================
    # 环境重置方法
    # ========================================================================

    def reset(self):
        """
        环境重置函数
        在每个episode开始时执行，负责：
        1. 向UE发送重置命令
        2. 接收UE返回的智能体属性信息
        3. 创建和配置智能体
        4. 获取初始观测数据
        5. 设置游戏配置（最大episode数等）

        Returns:
            初始观测数据
        """
        # 调用父类重置方法
        super().reset()
        self.t = 0  # 重置时间步
        pos_ro = np.random.rand() * 2 * np.pi  # 随机位置旋转角度

        # ====================================================================
        # 第一阶段：发送基础重置信息并接收UE端的智能体属性
        # ====================================================================

        resp = self.client.send_and_wait_reply(json.dumps({
            'valid': True,
            'DataCmd': 'reset_property',  # 重置属性命令
            'NumAgents': 0,
            'TimeStepMax': ScenarioConfig.MaxStepsPerEpisode,
            'TimeDilation': ScenarioConfig.TimeDilation,
            'FrameRate': ScenarioConfig.FrameRate,
            'TimeStep': 0,
            'Actions': None,
        }))

        resp = json.loads(resp)

        # ====================================================================
        # 第二阶段：解析智能体属性并创建智能体对象
        # ====================================================================

        n_team_agent = {}            # 统计每个队伍的智能体数量
        self.uid_to_agent_dict = {}  # 清空UID映射字典
        agent_property_list = resp['propertyArr']  # 从UE获取的智能体属性列表

        # 保存所有智能体属性（包括指挥官和中立方）
        self.AgentSettingAll = resp['propertyArr'].copy()
        index = 0
        self.AgentSettingArray.clear()

        # 遍历所有智能体属性，筛选并配置活跃智能体
        for i, agent_info in enumerate(agent_property_list):
            # 跳过指挥官类型的智能体
            if agent_info['type'] == 'BaseCommander_C':
                continue

            team = agent_info['agentTeam']
            # 跳过中立方（team != 0 且 team != 1）
            if team != 0 and team != 1:
                continue

            # 统计每个队伍的智能体数量
            if team not in n_team_agent:
                n_team_agent[team] = 0

            # 添加到活跃智能体数组
            self.AgentSettingArray.append(agent_info)

            # 配置智能体基本信息
            self.AgentSettingArray[index]['uid'] = agent_info['uId']
            self.AgentSettingArray[index]['team'] = team
            self.AgentSettingArray[index]['tid'] = n_team_agent[team]  # 队内ID
            self.AgentSettingArray[index]['type'] = agent_info['type']
            self.AgentSettingArray[index]['n_team_agent'] = n_team_agent[team]

            n_team_agent[team] += 1
            index += 1

        self.n_team_agent = n_team_agent

        # 2025/07 添加：配置每个队伍的智能体数量
        self.num_agent_each_team = [self.n_team_agent[i] for i in range(len(self.n_team_agent))]

        # 按队伍和队内ID排序，确保智能体顺序一致
        self.AgentSettingArray.sort(key=lambda x: (x['team'], x['tid']))

        # 更新全局配置
        self.SubTaskConfig.agent_list = self.AgentSettingArray.copy()
        GlobalConfig.ScenarioConfig.SubTaskConfig.agent_list = self.AgentSettingArray.copy()

        # 更新智能体总数
        self.n_agents = len(self.AgentSettingArray)
        self.agents.clear()

        # 创建智能体对象
        self.agents = [Agent(
            team=a['team'],
            team_id=a['tid'],
            uid=a['uid'],
            location=(a['initLocation']['x'], a['initLocation']['y'], a['initLocation']['z']),
            rotation=(a['initRotator']['pitch'], a['initRotator']['roll'], a['initRotator']['yaw']),
            velocity=(a['initVelocity']['x'], a['initVelocity']['y'], a['initVelocity']['z']),
            health=a['agentHp'],
            yaw=a['initRotator']['yaw'],
            speed=a['maxMoveSpeed']
        ) for a in self.AgentSettingArray]

        # ====================================================================
        # 第三阶段：发送重置观测命令并获取初始观测数据
        # ====================================================================

        resp = self.client.send_and_wait_reply(json.dumps({
            'valid': True,
            'DataCmd': 'reset_obs',  # 重置观测命令
            'NumAgents': 0,
            'TimeStepMax': ScenarioConfig.MaxStepsPerEpisode,
            'TimeDilation': ScenarioConfig.TimeDilation,
            'FrameRate': ScenarioConfig.FrameRate,
            'TimeStep': 0,
            'Actions': None,
        }))

        resp = json.loads(resp)

        # 验证所有智能体都已成功生成
        assert len(resp['dataArr']) == len(self.AgentSettingArray), \
            "Illegal agent initial position. 非法的智能体初始化位置，一部分智能体没有生成."

        # ====================================================================
        # 第四阶段：排序观测数据并解析全局配置
        # ====================================================================

        # 构建目标UID列表和观测字典，确保观测数据与智能体顺序一致
        target_uids = [agent['uid'] for agent in self.AgentSettingArray]
        obs_dict = {item['uId']: item for item in resp['dataArr']}
        sorted_obs = [obs_dict[uid] for uid in target_uids]
        resp['dataArr'] = sorted_obs

        # 解析全局配置信息（最大episode数、训练模式等）
        assert 'dataGlobal' in resp or 'rSVD1' in resp['dataGloble'], \
            f"The key dataGloble or rSVD1 does not exist in dictionary resp"

        config_str = resp['dataGlobal']['rSVD1']
        episode, train_mode = config_str.split(';')
        self.max_episode = int(episode)
        self.train_mode = bool(int(train_mode))

        # 打印当前episode信息并递增计数
        print('episode: %d' % (self.episode + 1))
        self.episode += 1
        # 个人修改查看参数回传
        print('sorted obs: %s' % str(sorted_obs))
        # 解析并返回观测信息
        return self.parse_response_ob_info(resp)

    # ========================================================================
    # 仿真步骤执行方法
    # ========================================================================

    def step(self, act):
        """
        仿真步骤执行函数
        负责执行单个仿真步骤，包括：
        1. 动作编码和适配
        2. 与UE通信执行动作
        3. 接收环境响应数据
        4. 计算奖励和判断任务结束
        5. 返回观测、奖励、完成标志和额外信息

        Args:
            act: 所有智能体的动作列表

        Returns:
            tuple: (观测, 奖励, 完成标志, 信息字典)
        """
        # 验证动作数量与智能体数量一致
        assert len(act) == self.n_agents

        # ====================================================================
        # 第一步：动作编码和适配
        # ====================================================================

        # 根据动作格式将数字动作转换为UE可识别的字符串动作
        if self.SubTaskConfig.ActionFormat == 'Single-Digit':
            act_send = [digit2act_dictionary[a] for a in act]
        elif self.SubTaskConfig.ActionFormat == 'Multi-Digit':
            act_send = [decode_action_as_string(a) for a in act]
        elif self.SubTaskConfig.ActionFormat == 'ASCII':
            act_send = [digitsToStrAction(a) for a in act]
        else:
            act_send = [digitsToStrAction(a) for a in act]

        # 动作适配：处理攻击动作的目标信息
        act_send = self.action_adaption(act_send)

        # ====================================================================
        # 第二步：与UE通信执行动作
        # ====================================================================

        try:
            resp = json.loads(self.client.send_and_wait_reply(json.dumps({
                'valid': True,
                'DataCmd': 'step',
                'TimeStep': self.t,
                'Actions': None,
                'StringActions': act_send,  # 字符串格式的动作
                'RSVD1': 'None'            # 保留字段，2025/05/29添加
            })))
        except TypeError:
            # 捕获通信错误
            return None
        except Exception as e:
            # 捕获其他异常
            return None

        # ====================================================================
        # 第三步：处理响应数据
        # ====================================================================

        # 排序观测数据，确保与智能体顺序一致
        target_uids = [agent['uid'] for agent in self.AgentSettingArray]
        obs_dict = {item['uId']: item for item in resp['dataArr']}
        sorted_obs = [obs_dict[uid] for uid in target_uids]
        resp['dataArr'] = sorted_obs

        # 更新当前步骤
        self.current_step = resp['dataGlobal']['timeCnt']

        # ====================================================================
        # 第四步：解析观测和计算奖励
        # ====================================================================

        # 获取观测数据和信息
        ob, info = self.parse_response_ob_info(resp)

        # 计算奖励和判断任务结束
        RewardForAllTeams, WinningResult = self.reward_and_done(resp)

        if WinningResult is not None:
            info.update(WinningResult)
            assert resp['dataGlobal']['episodeDone']
            done = True
        else:
            done = False

        # 验证最大步数限制
        if resp['dataGlobal']['timeCnt'] >= ScenarioConfig.MaxStepsPerEpisode:
            assert done

        return (ob, RewardForAllTeams, done, info)

    def action_adaption(self, action_list):
        """
        动作适配函数
        负责处理攻击动作，将目标从队伍ID转换为具体的目标UID

        Args:
            action_list: 原始动作列表

        Returns:
            list: 适配后的动作列表
        """
        for i, action in enumerate(action_list):
            agent_uid = self.agents[i].uid
            main_cmd, *rest = action.split(';')

            # 处理普通攻击动作
            if main_cmd.split('::')[1] == 'NormalAttacking':
                team, tid = rest
                # 查找目标智能体
                for j, agent in enumerate(self.agents):
                    if agent.team == int(team) and agent.team_id == int(tid):
                        # 将队伍信息替换为具体的目标UID
                        action_temp = 'ActionSet2::NormalAttacking;' + str(agent.uid) + ';' + str(agent_uid)
                        action_list[i] = action_temp
            else:
                # 为其他动作添加执行者UID
                action_list[i] = action + ';' + str(agent_uid)

        return action_list

    # ========================================================================
    # 数据解析和工具方法
    # ========================================================================

    def parse_event(self, event):
        """
        解析环境返回的关键事件
        如智能体阵亡、某队伍胜利等

        Args:
            event: 原始事件字符串

        Returns:
            dict: 解析后的事件字典
        """
        if not hasattr(self, 'pattern'):
            self.pattern = re.compile(r'<([^<>]*)>([^<>]*)')
        return {k: v for k, v in re.findall(self.pattern, event)}

    def extract_key_gameobj(self, resp):
        """
        获取非智能体的仿真物件，例如重要landmark等

        Args:
            resp: UE响应数据

        Returns:
            list: 关键游戏对象列表
        """
        keyObjArr = resp['dataGlobal']['keyObjArr']
        return keyObjArr

    def reward_and_done(self, resp):
        """
        奖励设计和任务结束判断
        建议在UE端定义触发奖励的事件，如智能体阵亡、战术目标完成等

        Args:
            resp: UE响应数据

        Returns:
            tuple: (奖励数组, 胜利结果)
        """
        # 初始化奖励数组
        reward = np.zeros(shape=(self.n_teams, self.n_agents))
        if self.Enable_Reward_Sharing:
            reward = np.zeros(shape=self.n_teams)

        # 初始化胜利结果
        WinningResult = {
            "team_ranking": [-1, -1],
            "end_reason": ''
        }

        return reward, WinningResult

    def step_skip(self):
        """
        跳过一次决策，暂无实际用途
        """
        return self.client.send_and_wait_reply(json.dumps({
            'valid': True,
            'DataCmd': 'skip_frame',
        }))

    def find_agent_by_uid(self, uid):
        """
        使用UID查找智能体（带缓存加速机制）

        Args:
            uid: 智能体唯一标识符

        Returns:
            Agent: 对应的智能体对象
        """
        # 更新UID到智能体的映射字典
        self.uid_to_agent_dict.update({agent.uid: agent for agent in self.agents})
        if isinstance(uid, str):
            self.uid_to_agent_dict.update({str(agent.uid): agent for agent in self.agents})

        return self.uid_to_agent_dict[uid]

    def parse_response_ob_info(self, resp):
        """
        解析智能体观测信息
        功能包括：
        1. 将死智能体的位置替换为无穷远
        2. 将智能体位置从字典形式转为简洁的(x,y,z)元组形式
        3. 更新智能体属性
        4. 提取关键游戏对象
        5. 配置观测参数

        Args:
            resp: UE响应数据

        Returns:
            tuple: (观测数据, 信息字典)
        """
        assert resp['valid']

        # 解析距离矩阵
        resp['dataGlobal']['distanceMat'] = np.array(
            resp['dataGlobal']['distanceMat']['flat_arr']
        ).reshape(self.n_agents, self.n_agents)

        info_dict = resp

        # 处理每个智能体的观测信息
        for info in info_dict['dataArr']:
            alive = info['agentAlive']

            if alive:
                # 存活智能体：解析位置、旋转、速度、缩放信息
                agentLocation = info.pop('agentLocation')
                agentRotation = info.pop('agentRotation')
                agentVelocity = info.pop('agentVelocity')
                agentScale = info.pop('agentScale')

                info['agentLocationArr'] = (agentLocation['x'], agentLocation['y'], agentLocation['z'])
                info['agentVelocityArr'] = (agentVelocity['x'], agentVelocity['y'], agentVelocity['z'])
                info['agentRotationArr'] = (agentRotation['yaw'], agentRotation['pitch'], agentRotation['roll'])
                info['agentScaleArr'] = (agentScale['x'], agentScale['y'], agentScale['z'])

                # 移除不需要的信息
                info.pop('previousAction')
                info.pop('availActions')
                info.pop('interaction')
            else:
                # 死亡智能体：位置设为无穷远
                inf = float('inf')
                info['agentLocationArr'] = (inf, inf, inf)
                info['agentVelocityArr'] = (inf, inf, inf)
                info['agentRotationArr'] = (inf, inf, inf)

        # 更新智能体属性
        info = resp['dataArr']
        for i, agent_info in enumerate(info):
            self.agents[i].update_agent_attrs(agent_info)

        # 提取关键游戏对象
        self.key_obj = self.extract_key_gameobj(resp)

        # 2025/07/08 根据智能体属性列表配置观测参数
        GlobalConfig.ScenarioConfig.SubTaskConfig.obs_n_entity = len(self.agents) + len(self.key_obj)
        GlobalConfig.ScenarioConfig.SubTaskConfig.keyobj_n_entity = len(self.key_obj)

        # 添加额外信息到信息字典
        info_dict.update({"AgentProperty": self.AgentSettingArray})
        info_dict.update({"AgentPropertyAll": self.AgentSettingAll})
        info_dict.update({"KeyObjNumber": len(self.key_obj)})
        info_dict.update({"TrainMode": self.train_mode})

        return self.make_obs(resp), info_dict

    # ========================================================================
    # 静态工具方法
    # ========================================================================

    @staticmethod
    def item_random_mv(src, dst, prob, rand=False):
        """
        随机移动算子
        可能用于实现随机打乱观测对象的顺序

        Args:
            src: 源数组
            dst: 目标数组
            prob: 移动概率
            rand: 是否随机打乱

        Returns:
            tuple: (处理后的源数组, 处理后的目标数组)
        """
        assert len(src.shape) == 1
        assert len(dst.shape) == 1
        if rand:
            np.random.shuffle(src)
        len_src = len(src)
        n_mv = (np.random.rand(len_src) < prob).sum()
        item_mv = src[range(len_src - n_mv, len_src)]
        src = src[range(0, 0 + len_src - n_mv)]
        dst = np.concatenate((item_mv, dst))
        return src, dst

    @staticmethod
    def get_binary_array(n_int, n_bits=8, dtype=np.float32):
        """
        工具类：将整数转换为二进制数组
        用于将智能体UID等信息编码为二进制表示

        Args:
            n_int: 输入整数数组
            n_bits: 二进制位数（默认8位）
            dtype: 数据类型

        Returns:
            np.ndarray: 二进制编码数组
        """
        arr = np.zeros((*n_int.shape, n_bits), dtype=dtype)
        for i in range(n_bits):
            arr[:, i] = (n_int % 2 == 1).astype(int)
            n_int = n_int / 2
            n_int = n_int.astype(np.int8)
        return arr

    # ========================================================================
    # 观测空间构建方法
    # ========================================================================

    def make_obs(self, resp=None, get_shape=False):
        """
        构建观测空间
        这是环境观测的核心方法，负责：
        1. 计算智能体间的距离矩阵
        2. 选择每个智能体附近的友方和敌方单位
        3. 构建包含智能体状态、位置、速度等信息的观测向量
        4. 添加关键游戏对象信息

        Args:
            resp: UE响应数据
            get_shape: 是否只返回观测空间形状

        Returns:
            观测数据或观测空间形状
        """
        CORE_DIM = 23  # 核心观测维度
        assert ScenarioConfig.obs_vec_length == CORE_DIM

        if get_shape:
            return CORE_DIM

        # 临时参数
        OBS_RANGE_PYTHON_SIDE = 8000  # Python端观测范围
        n_key_obj = self.SubTaskConfig.keyobj_n_entity  # 关键对象数量

        # 计算观测中敌方和友方单位的最大数量
        MAX_NUM_OPP_OBS = (self.SubTaskConfig.obs_n_entity - n_key_obj) // 2
        MAX_NUM_ALL_OBS = (self.SubTaskConfig.obs_n_entity - n_key_obj) // 2

        # ====================================================================
        # 第一步：计算距离矩阵
        # ====================================================================

        # 获取所有智能体的3D位置
        pos3d_arr = np.zeros(shape=(self.n_agents, 3), dtype=np.float32)
        for i, agent in enumerate(self.agents):
            pos3d_arr[i] = agent.pos3d

        # 使用UE计算的距离矩阵加速
        dis_mat = resp['dataGlobal']['distanceMat']
        alive_all = np.array([agent.alive for agent in self.agents])

        # 死亡智能体的距离设为无穷远
        dis_mat[~alive_all, :] = +np.inf
        dis_mat[:, ~alive_all] = +np.inf

        # 获取队伍归属信息
        team_belonging = np.array([agent.team for agent in self.agents])

        # ====================================================================
        # 第二步：构建所有智能体的基础观测
        # ====================================================================

        obs_arr = RawObsArray(key='Agent')

        # 生成智能体UID的二进制编码（缓存机制）
        if not hasattr(self, "uid_binary"):
            self.uid_binary = self.get_binary_array(np.arange(self.n_agents), 10)

        # 为每个智能体构建观测向量
        for i, agent in enumerate(self.agents):
            assert agent.location is not None, ("agent name:%s" % agent.uid)

            # 添加UID二进制编码 (0-9位)
            obs_arr.append(self.uid_binary[i])

            # 添加基础信息 (10-13位)
            obs_arr.append([
                agent.index,     # 10: 智能体索引
                agent.team,      # 11: 队伍
                agent.alive,     # 12: 存活状态
                agent.uid_remote, # 13: 远程UID
            ])

            # 添加3D位置信息 (14-16位)
            obs_arr.append(agent.pos3d)

            # 添加3D速度信息 (17-19位)
            obs_arr.append(agent.vel3d)

            # 添加其他状态信息 (20-22位)
            obs_arr.append([
                agent.hp,        # 20: 生命值
                agent.yaw,       # 21: 偏航角
                agent.max_speed, # 22: 最大速度
            ])

        obs_ = obs_arr.get()
        new_obs = my_view(obs_, [self.n_agents, -1])
        assert CORE_DIM == new_obs.shape[-1]

        # ====================================================================
        # 第三步：为每个智能体构建局部观测
        # ====================================================================

        # 初始化观测数组：形状为(智能体数, 观测对象数, 观测维度)
        OBS_ALL_AGENTS = np.zeros(shape=(
            self.n_agents,
            MAX_NUM_OPP_OBS + MAX_NUM_ALL_OBS,
            CORE_DIM
        ))

        # 为每个智能体构建局部观测
        for i, agent in enumerate(self.agents):
            if not agent.alive:
                OBS_ALL_AGENTS[i, :] = np.nan  # 死亡智能体的观测设为NaN
                continue

            # 计算到所有智能体的距离
            dis2all = dis_mat[i, :]
            is_ally = (team_belonging == agent.team)

            # =================================================================
            # 处理敌方观测
            # =================================================================

            a2h_dis = dis2all[~is_ally]      # 与敌方的距离
            h_alive = alive_all[~is_ally]    # 敌方存活状态
            h_feature = new_obs[~is_ally]    # 敌方观测特征

            # 按距离排序，选择最近的敌方单位
            h_iden_sort = np.argsort(a2h_dis)[:MAX_NUM_OPP_OBS]
            a2h_dis_sorted = a2h_dis[h_iden_sort]
            h_alive_sorted = h_alive[h_iden_sort]
            h_vis_mask = (a2h_dis_sorted <= OBS_RANGE_PYTHON_SIDE) & h_alive_sorted

            # 分离可见和不可见的敌方单位
            h_vis_index = h_iden_sort[h_vis_mask]
            h_invis_index = h_iden_sort[~h_vis_mask]

            # 随机打乱可见敌方单位
            h_vis_index, h_invis_index = self.item_random_mv(
                src=h_vis_index, dst=h_invis_index, prob=0, rand=True
            )

            # 拼接索引并构建掩码
            h_ind = np.concatenate((h_vis_index, h_invis_index))
            h_msk = np.concatenate((h_vis_index < 0, h_invis_index >= 0))

            # 获取敌方观测特征并应用掩码
            a2h_feature_sort = h_feature[h_ind]
            a2h_feature_sort[h_msk] = 0  # 不可见的敌方信息设为0

            # 填充不足的观测槽位
            if len(a2h_feature_sort) < MAX_NUM_OPP_OBS:
                a2h_feature_sort = np.concatenate((
                    a2h_feature_sort,
                    np.zeros(shape=(MAX_NUM_OPP_OBS - len(a2h_feature_sort), CORE_DIM))
                ), axis=0)

            # =================================================================
            # 处理友方观测（包括自身）
            # =================================================================

            a2f_dis = dis2all[is_ally]      # 与友方的距离
            f_alive = alive_all[is_ally]    # 友方存活状态
            f_feature = new_obs[is_ally]    # 友方观测特征

            # 按距离排序，选择最近的友方单位
            f_iden_sort = np.argsort(a2f_dis)[:MAX_NUM_ALL_OBS]
            a2f_dis_sorted = a2f_dis[f_iden_sort]
            f_alive_sorted = f_alive[f_iden_sort]
            f_vis_mask = (a2f_dis_sorted <= OBS_RANGE_PYTHON_SIDE) & f_alive_sorted

            # 分离自身、可见友方和不可见友方
            f_vis_index = f_iden_sort[f_vis_mask]
            self_vis_index = f_vis_index[:1]    # 自身
            f_vis_index = f_vis_index[1:]       # 其他友方
            f_invis_index = f_iden_sort[~f_vis_mask]

            # 随机打乱可见友方
            f_vis_index, f_invis_index = self.item_random_mv(
                src=f_vis_index, dst=f_invis_index, prob=0, rand=True
            )

            # 拼接索引并构建掩码
            f_ind = np.concatenate((self_vis_index, f_vis_index, f_invis_index))
            f_msk = np.concatenate((self_vis_index < 0, f_vis_index < 0, f_invis_index >= 0))

            # 获取友方观测特征并应用掩码
            self_ally_feature_sort = f_feature[f_ind]
            self_ally_feature_sort[f_msk] = 0  # 不可见的友方信息设为0

            # 填充不足的观测槽位
            if len(self_ally_feature_sort) < MAX_NUM_ALL_OBS:
                self_ally_feature_sort = np.concatenate((
                    self_ally_feature_sort,
                    np.zeros(shape=(MAX_NUM_ALL_OBS - len(self_ally_feature_sort), CORE_DIM))
                ), axis=0)

            # =================================================================
            # 组合友方和敌方观测
            # =================================================================

            OBS_ALL_AGENTS[i, :] = np.concatenate(
                (self_ally_feature_sort, a2h_feature_sort), axis=0
            )

        # ====================================================================
        # 第四步：添加关键游戏对象观测
        # ====================================================================

        MAX_OBJ_NUM_ACCEPT = 1  # 最多接受的关键对象数量
        self.N_Obj = len(self.key_obj)
        OBJ_UID_OFFSET = 32768  # 对象UID偏移量

        # 构建关键对象观测
        obs_arr = RawObsArray(key='GameObj')

        for i, obj in enumerate(self.key_obj):
            assert obj['uId'] - OBJ_UID_OFFSET == i

            # 添加反向UID二进制编码（区分智能体和对象）
            obs_arr.append(-self.uid_binary[i])

            # 添加对象基础信息
            obs_arr.append([
                obj['uId'] - OBJ_UID_OFFSET,  # 对象索引
                -1,                           # 队伍标记（-1表示非智能体）
                True,                         # 存活状态
                obj['uId'] - OBJ_UID_OFFSET,  # 对象UID
            ])

            # 添加对象位置信息
            obs_arr.append([
                obj['location']['x'],
                obj['location']['y'],
                obj['location']['z']
            ])

            # 添加对象速度和其他属性
            obs_arr.append([
                obj['velocity']['x'],
                obj['velocity']['y'],
                obj['velocity']['z'],
                -1,                           # 生命值（-1表示不适用）
                obj['rotation']['yaw'],       # 偏航角
                0,                            # 最大速度（0表示不适用）
            ])

        # 处理关键对象观测
        OBS_GameObj = my_view(obs_arr.get(), [len(self.key_obj), -1])[:MAX_OBJ_NUM_ACCEPT, :]
        OBS_GameObj = repeat_at(OBS_GameObj, insert_dim=0, n_times=self.n_agents)

        # 将关键对象观测合并到总观测中
        OBS_ALL_AGENTS = np.concatenate((OBS_ALL_AGENTS, OBS_GameObj), axis=1)

        return OBS_ALL_AGENTS

    # ========================================================================
    # 智能体初始化方法
    # ========================================================================

    def init_ground(self, agent_info, pos_ro):
        """
        初始化地面智能体
        在网格状排列中创建地面单位

        Args:
            agent_info: 智能体信息字典
            pos_ro: 位置旋转角度

        Returns:
            dict: 智能体属性配置
        """
        N_COL = 4  # 网格列数
        agent_class = agent_info['type']
        team = agent_info['team']
        n_team_agent = agent_info['n_team_agent']
        tid = agent_info['tid']
        uid = agent_info['uid']

        # 计算初始位置（网格排列）
        x = 0 + 800 * (tid - n_team_agent // 2) // N_COL
        y = (400 * (tid % N_COL) + 2000) * (-1) ** (team + 1)

        # 应用旋转
        x, y = np.matmul(np.array([x, y]),
                        np.array([[np.cos(pos_ro), -np.sin(pos_ro)],
                                 [np.sin(pos_ro), np.cos(pos_ro)]]))
        z = 500  # 略高于地面

        # 设置初始朝向
        yaw = 90 if team == 0 else -90

        # 验证位置在地图范围内
        assert np.abs(x) < 15000.0 and np.abs(y) < 15000.0

        # 创建智能体属性配置
        agent_property = copy.deepcopy(AgentPropertyDefaults)
        agent_property.update({
            'DebugAgent': False,
            # 最大移动速度
            'MaxMoveSpeed': 720 if agent_class == 'RLA_CAR_Laser' else 600,
            # 智能体缩放（影响质量，谨慎修改）
            'AgentScale': {'x': 0.5, 'y': 0.5, 'z': 0.5},
            # 闪避概率
            "DodgeProb": 0.0,
            # 爆炸伤害
            "ExplodeDmg": 20,
            # 队伍归属
            'AgentTeam': team,
            # UE类名
            'ClassName': agent_class,
            # 武器冷却时间
            'WeaponCD': 1,
            # 感知范围、警戒范围、开火范围
            "PerceptionRange": 2000 if agent_class == 'RLA_CAR_Laser' else 2500,
            "GuardRange": 1400 if agent_class == 'RLA_CAR_Laser' else 1700,
            "FireRange": 750 if agent_class == 'RLA_CAR_Laser' else 1400,
            # 调试信息：感知环范围
            'RSVD1': '-Ring1=2000 -Ring2=1400 -Ring3=750' if agent_class == 'RLA_CAR_Laser' else '-Ring1=2500 -Ring2=1700 -Ring3=1400',
            # 初始动作设置
            'RSVD2': '-InitAct=ActionSet2::Idle;AsFarAsPossible',
            # 智能体生命值（随机化）
            'AgentHp': np.random.randint(low=95, high=105) if agent_class == 'RLA_CAR_Laser' else np.random.randint(low=145, high=155),
            # 队内排名
            'IndexInTeam': tid,
            # 系统唯一标识符
            'UID': uid,
            # 显示颜色
            'Color': '(R=0,G=1,B=0,A=1)' if team == 0 else '(R=0,G=0,B=1,A=1)',
            # 初始位置
            'InitLocation': {'x': x, 'y': y, 'z': z},
            # 初始旋转
            'InitRotator': {'pitch': 0, 'roll': 0, 'yaw': yaw},
        })

        return agent_property

    def init_location(self, agent_info, pos_ro):
        """
        使用指定位置初始化智能体
        根据agent_info中提供的坐标创建智能体

        Args:
            agent_info: 包含位置信息的智能体信息字典
            pos_ro: 位置旋转角度（此处未使用）

        Returns:
            dict: 智能体属性配置
        """
        agent_class = agent_info['type']
        team = agent_info['team']
        n_team_agent = agent_info['n_team_agent']
        tid = agent_info['tid']
        uid = agent_info['uid']

        # 使用指定的坐标
        x = agent_info['x']
        y = agent_info['y']
        z = agent_info['z']
        yaw = 90 if team == 0 else -90  # 队伍决定朝向

        # 创建智能体属性配置
        agent_property = copy.deepcopy(AgentPropertyDefaults)
        agent_property.update({
            'DebugAgent': False,
            # 最大移动速度
            'MaxMoveSpeed': 720 if agent_class == 'RLA_CAR_Laser' else 600,
            # 智能体缩放
            'AgentScale': {'x': 0.8, 'y': 0.8, 'z': 0.8},
            # 闪避概率
            "DodgeProb": 0.0,
            # 爆炸伤害
            "ExplodeDmg": 10,
            # 队伍归属
            'AgentTeam': team,
            # UE类名
            'ClassName': agent_class,
            # 武器冷却时间
            'WeaponCD': 2,
            # 扩展的感知和战斗范围
            "PerceptionRange": 5000 if agent_class == 'RLA_CAR_Laser' else 5500,
            "GuardRange": 2800 if agent_class == 'RLA_CAR_Laser' else 3400,
            "FireRange": 3500 if agent_class == 'RLA_CAR_Laser' else 4800,
            # 调试信息
            'RSVD1': '-Ring1=2000 -Ring2=1400 -Ring3=750' if agent_class == 'RLA_CAR_Laser' else '-Ring1=2500 -Ring2=1700 -Ring3=1400',
            # 初始动作设置
            'RSVD2': '-InitAct=ActionSet2::Idle;AsFarAsPossible',
            # 固定生命值（不随机化）
            'AgentHp': 105 if agent_class == 'RLA_CAR_Laser' else 155,
            # 队内排名
            'IndexInTeam': tid,
            # 系统唯一标识符
            'UID': uid,
            # 团队奖励标志
            "IsTeamReward": False,
            # 显示颜色（不同队伍不同颜色）
            'Color': '(R=1,G=0,B=0,A=1)' if team == 0 else '(R=0,G=0,B=1,A=1)',
            # 初始位置
            'InitLocation': {'x': x, 'y': y, 'z': z},
            # 初始旋转
            'InitRotator': {'pitch': 0, 'roll': 0, 'yaw': yaw},
        })

        return agent_property

    def init_air(self, agent_info, pos_ro):
        """
        初始化空中智能体
        在空中创建飞行单位

        Args:
            agent_info: 智能体信息字典
            pos_ro: 位置旋转角度

        Returns:
            dict: 智能体属性配置
        """
        N_COL = 4  # 网格列数
        agent_class = agent_info['type']
        team = agent_info['team']
        n_team_agent = agent_info['n_team_agent']
        tid = agent_info['tid']
        uid = agent_info['uid']

        # 计算空中初始位置
        x = 0 + 800 * (tid - n_team_agent // 2) // N_COL
        y = 2000 * (-1) ** (team + 1)

        # 应用旋转
        x, y = np.matmul(np.array([x, y]),
                        np.array([[np.cos(pos_ro), -np.sin(pos_ro)],
                                 [np.sin(pos_ro), np.cos(pos_ro)]]))
        z = 1000  # 空中高度

        # 设置初始朝向
        yaw = 90 if team == 0 else -90

        # 验证位置在地图范围内
        assert np.abs(x) < 15000.0 and np.abs(y) < 15000.0

        # 创建空中智能体属性配置
        agent_property = copy.deepcopy(AgentPropertyDefaults)
        agent_property.update({
            'DebugAgent': False,
            # 更高的空中移动速度
            'MaxMoveSpeed': 900,
            # 空中智能体缩放
            'AgentScale': {'x': 0.5, 'y': 0.5, 'z': 0.5},
            # 闪避概率
            "DodgeProb": 0.0,
            # 爆炸伤害
            "ExplodeDmg": 10,
            # 队伍归属
            'AgentTeam': team,
            # UE类名
            'ClassName': agent_class,
            # 武器冷却时间（空中单位冷却时间更长）
            'WeaponCD': 3,
            # 空中智能体的感知和战斗范围
            "PerceptionRange": 2500,
            "GuardRange": 1800,
            "FireRange": 1700,
            # 调试信息
            'RSVD1': '-ring1=2500 -ring2=1800 -ring3=1700',
            # 初始动作设置（空中单位使用StaticAlert）
            'RSVD2': '-InitAct=ActionSet2::Idle;StaticAlert',
            # 空中智能体生命值（较低）
            'AgentHp': 50,
            # 队内排名
            'IndexInTeam': tid,
            # 系统唯一标识符
            'UID': uid,
            # 显示颜色
            'Color': '(R=0,G=1,B=0,A=1)' if team == 0 else '(R=0,G=0,B=1,A=1)',
            # 初始位置
            'InitLocation': {'x': x, 'y': y, 'z': z},
            # 初始旋转
            'InitRotator': {'pitch': 0, 'roll': 0, 'yaw': yaw},
        })

        return agent_property

# ============================================================================
# 文件总结：
#
# CssimAgentEnvBase.py是CSSIM仿真系统的核心环境类，提供了：
# 1. 智能体生命周期管理（创建、重置、更新）
# 2. 与Unreal Engine的双向通信
# 3. 复杂的观测空间构建（包括局部观测选择）
# 4. 动作编码、适配和执行
# 5. 奖励计算框架
# 6. 多种智能体初始化策略（地面、空中、指定位置）
#
# 该类是强化学习训练的基础环境接口，支持多智能体、多队伍的战术仿真。
# ============================================================================